import numpy as np
import pandas as pd

import util.mydecorator as mydecorator
import util.lib as lib

import pandas as pd

from format_drawdown import job as format_job

from daily import get_daily_df

def get_stocks():
    return [
        "纳指ETF",
        "黄金ETF",
        "可转债ETF",
        "红利ETF",
        "豆粕ETF",
        "标普500ETF",
        "30年期国债ETF",
        "上证指数",
        "中概互联ETF",
        "工商银行",
        "招商银行",
        "长江电力",
        "国投白银ETF",
        "南方原油LOF",
        "上证指数ETF",
        "科创50ETF",
        "煤炭ETF",
        "电池ETF",
        "电力ETF",
        "银行ETF",
        "基建ETF",
        "家电ETF",
        "半导体ETF",
        "证券保险ETF",
        "五粮液",
        "美的集团",
        "招商公路",
    ]
def calculate_drawdown(writer, stock):
    
    df = get_daily_df(stock).tail(2300).copy()
    
    prices = df["close"]

    # 计算累计最大值（峰值）
    cumulative_max = prices.cummax()
    
    # 计算每日回撤（从峰值的下跌百分比）
    drawdown = round( 100 * (prices / cumulative_max - 1), 2)
    df["drawdown"] = drawdown
    df["stock"] = stock

    df = df.sort_index(ascending=False)
    df[["stock","close",
        "drawdown",
        "pct",'pct_3d','pct_5d',
        "level_point","rel_lp",
        "gt30d","gt60d","gt100d","gt300d",
        "open_pct","pct","low_pct","high_pct","next_pct"
        ]].to_excel(writer, sheet_name=stock)


@mydecorator.calTime
def job():
    writer = pd.ExcelWriter(lib.data_path + "drawdown.xlsx")

    stocks = get_stocks()
    for stock in stocks:
        print(stock)
        calculate_drawdown(writer, stock)

    writer.close()

# 回撤统计
if __name__ == "__main__":
    job()
    format_job()